
import isaaclab_tasks.manager_based.locomotion.velocity.mdp as mdp
from isaaclab.managers import RewardTermCfg as RewTerm
from isaaclab.managers import SceneEntityCfg
from isaaclab.utils import configclass

from legged_lab.tasks.manager_based.locomotion.velocity_env_cfg import LocomotionVelocityRoughEnvCfg, RewardsCfg
from legged_lab.assets.config.sf_tron1 import SF_TRON1_CFG

##
# Pre-defined configs
##
from isaaclab_assets.robots.unitree import G1_MINIMAL_CFG  # isort: skip


@configclass
class SfTron1RewardsCfg(RewardsCfg):
    """Reward terms for the MDP."""
    feet_air_time = RewTerm(
        func=mdp.feet_air_time_positive_biped,
        weight=30.0,
        params={
            "command_name": "base_velocity",
            "sensor_cfg": SceneEntityCfg("contact_forces", body_names="ankle_.*"),
            "threshold": 0.2,
        },
    )


@configclass
class SfTron1RoughEnvCfg(LocomotionVelocityRoughEnvCfg):

    base_link_name = "base_Link"
    foot_link_name = "ankle_.*"

    rewards: SfTron1RewardsCfg = SfTron1RewardsCfg()

    def __post_init__(self):
        # post init of parent
        super().__post_init__()

        # ------------------------------Sence------------------------------
        # switch robot 
        self.scene.robot = SF_TRON1_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
        self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name

        # ------------------------------Observations------------------------------
        self.observations.policy.base_lin_vel = None
        self.observations.policy.height_scan = None

        # ------------------------------Actions------------------------------
        # reduce action scale
        # self.actions.joint_pos.scale = 1.0

        # ------------------------------Events------------------------------
        self.events.randomize_rigid_body_mass.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_rigid_body_mass.params["mass_distribution_params"] = (-1.0, 2.0)
        self.events.randomize_rigid_body_inertia.params["inertia_distribution_params"] = (0.9, 1.2)
        self.events.randomize_com_positions.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_com_positions.params["com_distribution_params"] = (-0.03, 0.03)
        self.events.randomize_apply_external_force_torque = None
        self.events.randomize_reset_joints.params["velocity_range"] = (0.0, 0.0)


        # ------------------------------Rewards------------------------------
        # General
        # self.rewards.is_terminated.weight = -200

        # Root penalties
        self.rewards.lin_vel_z_l2.weight = -0.5
        self.rewards.ang_vel_xy_l2.weight = -0.05
        self.rewards.flat_orientation_l2.weight = -5.0
        self.rewards.base_height_l2.params["target_height"] = 0.8
        self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]
        self.rewards.base_height_l2.weight = -10.0

        # Joint penaltie
        self.rewards.joint_torques_l2.weight = -2.5e-05
        # self.rewards.joint_vel_l2.weight = 0
        self.rewards.joint_acc_l2.weight = -2.5e-7
        # self.rewards.create_joint_deviation_l1_rewterm(
        #     "joint_deviation_other_l1", -0.2, [".*_hip_yaw", ".*_hip_roll", ".*_shoulder_.*", ".*_elbow"]
        # )
        # self.rewards.create_joint_deviation_l1_rewterm("joint_deviation_torso_l1", -0.1, ["torso"])
        # self.rewards.joint_pos_limits.weight = -1.0
        # self.rewards.joint_pos_limits.params["asset_cfg"].joint_names = [".*_ankle"]
        # self.rewards.joint_vel_limits.weight = 0
        # self.rewards.joint_power.weight = 0
        self.rewards.create_joint_deviation_l1_rewterm(
            "joint_deviation_other_l1", -0.2, ["abad_.*", "hip_.*", "knee_.*"]
        )
        self.rewards.create_joint_deviation_l1_rewterm(
            "joint_deviation_ankle_l1", -0.25, ["ankle_.*"]
        )
        self.rewards.stand_still_without_cmd.weight = -5.0
        self.rewards.joint_pos_penalty.weight = -1.0

        # Action penalties
        self.rewards.action_rate_l2.weight = -0.01

        # Contact sensor
        self.rewards.undesired_contacts.weight = -50.0
        self.rewards.undesired_contacts.params["sensor_cfg"].body_names=["base_.*", "abad_.*", "hip_.*"]
        # self.rewards.contact_forces.weight = 0

        # Velocity-tracking rewards
        self.rewards.track_lin_vel_xy_exp.weight = 15.0
        self.rewards.track_lin_vel_xy_exp.func = mdp.track_lin_vel_xy_yaw_frame_exp
        self.rewards.track_ang_vel_z_exp.weight = 7.5
        self.rewards.track_ang_vel_z_exp.func = mdp.track_ang_vel_z_world_exp

        # Others
        # self.rewards.feet_air_time.weight = 1.0
        # self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_air_time.params["threshold"] = 0.6
        # self.rewards.feet_contact.weight = 0.0
        self.rewards.feet_slide.weight = -0.25
        self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_distance_y_exp_for_biped.weight = 10.0
        self.rewards.feet_distance_y_exp_for_biped.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_distance_y_exp_for_biped.params["stance_width"] = 0.25 
        self.rewards.upward.weight = 1.0

        # If the weight of rewards is 0, set rewards to None
        if self.__class__.__name__ == "SfTron1RoughEnvCfg":
            self.disable_zero_weight_rewards()

        # ------------------------------Terminations------------------------------
        self.terminations.illegal_contact.params["sensor_cfg"].body_names = [self.base_link_name]

        # ------------------------------Commands------------------------------
        self.commands.base_velocity.ranges.lin_vel_x = (-1.0, 1.0)
        self.commands.base_velocity.ranges.lin_vel_y = (-0.8, 0.8)
        self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0)
